import os
from quant_researcher.quant.project_tool.file_tool import get_all_filename_path
from quant_researcher.quant.project_tool.time_tool import date_shifter, calc_date_diff, get_yesterday
from quant_researcher.quant.datasource_fetch.crypto_api.coinmarketcap import get_asset_ohlcvm_via_http
from quant_researcher.quant.project_tool.db_operator import db_conn, conn_dispatcher
import pandas as pd
import random
import math
import time


def split_a_list_into_batches(a_list, how_many):
    """
    把一个 list 分成几份，尽量均分

    :raise ValueError: 如果传入的份数 = 0
    :param a_list: 数据
    :param int how_many: 几份
    :return: list，每一个 element 都是源数据的子集
    """
    # 2019年12月22日14:22:21，如果这个数组为空的话怎么办？为了保持一致性
    # 因为这个函数的结果一般是用来进行循环的，所以返回一个空数组看起来是比较合理的
    if not a_list:
        print(f'传入的数组判断为False：{a_list}，为了保持结果的一致性，返回[]')
        return []

    if how_many == 0:
        raise ValueError('需要分成的进程数不能是 0')
    total_count = len(a_list)
    print(f'总元素个数：{total_count}，希望进程数：{how_many}')
    actual_pool_size = min(total_count, how_many)
    step_max = math.ceil(total_count / actual_pool_size)
    step_min = math.floor(total_count / actual_pool_size)
    batch_list_size = [step_min] * actual_pool_size
    gap = total_count - (step_min * actual_pool_size)
    batch_list_size[:gap] = [step_max] * gap
    print(f'{actual_pool_size}个进程, 每个进程分到的元素个数：{batch_list_size}')

    start_idx = 0
    batch_list = []
    for idx in batch_list_size:
        end_idx = start_idx + idx
        batch_list.append(a_list[start_idx:end_idx])
        start_idx = end_idx
    assert sum([len(x) for x in batch_list]) == total_count

    return batch_list


def batch_process(a_list, func, **kwargs):
    """
    批量处理，把 list 中的每个对象都交给 func 处理并附上指定的关键字参数 info； 这里的list是已经进行分割好的batch

    :param a_list: 其中每一个对象都会交给func处理
    :param func: 处理的函数
    :param kwargs:
        - multiple_threads，bool，是否使用多线程，默认：True
        - info，dict，将会作为关键字参数传递给func，默认：{}
        - step，int，如果是多线程处理，每次最大的线程数，默认：30
        - ignore_err: bool，是否忽略错误，默认：False
        - pool_size: int，目标函数执行的最多进程数，默认：None
        - verbose，bool，是不是把函数执行的第一个参数也一并使用 tuple 返回
            方便在多进程和多线程拿到结果的时候对比，要不然是乱的，默认：False
    :return:
        list，每一个对象时每次func运行的返回值
    """
    ignore_err = kwargs.pop('ignore_err', False)
    multiple_threads = kwargs.pop('multiple_threads', True)
    verbose = kwargs.pop('verbose', False)

    assert_tool.value_is_instance(a_list, list)

    total_count = len(a_list)
    LOG.info(f'输入数组元素个数={total_count}，verbose={verbose}，'
             f'是否多线程={multiple_threads}，处理函数：{func.__name__}')

    # 2019年12月22日14:22:21，如果这个数组为空的话怎么办？为了保持一致性
    # 因为这个函数的结果一般是用来进行循环的，所以返回一个空数组看起来是比较合理的
    if not a_list:
        LOG.warning(f'传入的数组判断为False：{a_list}，为了保持结果的一致性，返回[]')
        return []

    total_res = []
    if multiple_threads:
        LOG.info('将采用多线程来执行任务')
        step = kwargs.pop('step', 30)
        LOG.info('多线程时的最大线程同时数量: %s', step)

        def for_a_batch(batch, info):
            """deal with a batch"""
            res = []
            thread_list = []
            for a_code in batch:
                # 输入总是传给func的第一个参数
                a_thread = MultiThread(func, args=(a_code,), info=info)
                a_thread.start()
                thread_list.append((a_code, a_thread))
            for item in thread_list:
                a_thread = item[1]
                a_thread.join()
                a_res = a_thread.get_result()
                if verbose:
                    res.append((item[0], a_res))
                else:
                    res.append(a_res)
            return res

        start_idx = 0
        while start_idx < total_count:
            end_idx = min(start_idx + step, total_count)
            this_batch = a_list[start_idx:end_idx]
            total_res.extend(for_a_batch(this_batch, kwargs.get('info', {})))
            start_idx = end_idx
    else:
        pool_size = kwargs.pop('pool_size', None)
        if pool_size is None:
            raise ValueError('传入的 pool size 是 None，或者根本没传')
        pool_size = min(pool_size, total_count)
        LOG.info(f'将采用多进程来执行任务！'
                 f'任务数量：{total_count}，进程池大小：{pool_size}')
        pool = multiprocessing.Pool(pool_size)
        for the_one in a_list:
            # 理论上，Python是不能pickle函数的
            # 只能将函数名和模块名进行pickle，所以就不能给函数加wrapper
            # apply_async(
            #     func, args=(), kwds={}, callback=None, error_callback=None
            # )
            kwds = kwargs.get('info', {})
            kwds.update({'ignore_err': ignore_err})
            total_res.append(pool.apply_async(func, args=(the_one,), kwds=kwds))
        LOG.warning('所有数据已加入进程池，即将关闭进程池')
        pool.close()
        LOG.warning('已关闭进程池，等待所有进程结束 ...')
        pool.join()
        LOG.warning(f'所有进程都已结束！已获得结果数量：{len(total_res)}')
        total_res = [r.get() for r in total_res]

    assert len(total_res) == len(a_list)
    return total_res


def _deal_a_batch_with_a_loop_v2(a_batch, func, **kwargs):
    """
    循环将一个 list 中的元素传入指定函数运行，并将每一次运行的结果放入 list 返回； ; 这里是将所有任务分成多个batch后，对一个batch内的多任务进行循环处理

    与上述区别在于传入每个进程时，不需要提前创建好数据库连接，只需要制定数据库连接名
    这样的好处时，每个进程使用的数据库连接是分开的

    :param list a_batch: 数据
    :param func: 函数
    :param kwargs:
        - prepare_these_conn，list，需准备的 conn 的名字，默认：[]
        - ignore_err，bool，忽略某一次运行的发生的错误，默认：False
    :return: list
    """
    prepare_these_conn = kwargs.pop('prepare_these_conn', [])
    ignore_err = kwargs.pop('ignore_err', False)
    # LOG.info(f'需准备的连接：{prepare_these_conn}，是否忽略错误：{ignore_err}')

    conn_dict = None
    if prepare_these_conn:
        conn_dict = conn_dispatcher.dispatcher(db_names=prepare_these_conn)
        kwargs.update({'conn_dict': conn_dict})

    res = []
    for x in a_batch:
        try:
            res.append(func(x, **kwargs))
        except Exception as err:
            print(f'处理 {x} 时发生错误：{err}')
            if not ignore_err:
                conn_dispatcher.close_them(conn_dict)
                raise err
    conn_dispatcher.close_them(conn_dict)
    print(f'该进程任务处理结束')
    return res


def fetch_one_coin_data(id_series, **kwargs):
    end_date = kwargs.pop('end_date')
    file_path = kwargs.pop('file_path')
    initialize = kwargs.pop('initialize')
    gap_days = kwargs.pop('gap_days')  # 默认更新至最新两天
    all_existing_coins_name = kwargs.pop('all_existing_coins_name')
    all_coin_list_df = kwargs.pop('all_coin_list_df')
    all_filename_path_dict = kwargs.pop('all_filename_path_dict')

    for id in id_series:
        print(f'开始获取id={id}的数据')
        coins = id_series[id_series == id].index[0]  # 币的名称
        if (coins in all_existing_coins_name):  # 如果这个币之前获取过
            file_name = all_filename_path_dict[f'{coins}.xlsx']
            history_ohlcvm = pd.read_excel(f'{file_name}', index_col='end_date')
            if history_ohlcvm.empty:
                history_ohlcvm_latest_date = '2013-01-01'
            else:
                history_ohlcvm_latest_date = history_ohlcvm.index[-1]
        else:  # 之前没获取过，则默认都需要进行获取
            file_name = os.path.join(file_path, f'{coins}.xlsx')
            history_ohlcvm_latest_date = '2013-01-01'

        if initialize:  # 如果是初始化的话，只要文件夹里已经包含这个文件就不进行更新, 因为中间可能中断过，文件夹里有的就不重新获取了
            mask = (coins in all_existing_coins_name)
        else:  # 如果是不是初始化的话，需要判断数据是不是最新，并且这个币是不是处于inactive状态
            mask1 = (calc_date_diff(history_ohlcvm_latest_date, end_date, fmt_str='%Y-%m-%d') < gap_days)  # 该币已经更新到最近2天了
            mask2 = bool(all_coin_list_df[all_coin_list_df['id'] == id]['is_active'][0])
            mask = mask1 | (not mask2)  # 数据已经更新到最近gap_days天了，或者数据是inactive 才不进行更新，否则需要进行更新

        if mask:  # 无需更新
            print(f'{coins} 数据已最新,或者数据已经不更新了，无需再更新')
            # all_ohlcvm = history_ohlcvm.copy()
            # all_ohlcvm.reset_index(inplace=True)
            # all_ohlcvm.drop_duplicates(subset=['end_date'], inplace=True)  # 同一天多条记录进行去重
            # all_ohlcvm.set_index('end_date', inplace=True)
            # all_ohlcvm.sort_index(inplace=True)
            # all_ohlcvm.index.name = 'end_date'
            continue
        else:  # 需要更新或者进行头一次获取
            all_ohlcvm = get_asset_ohlcvm_via_http(id=id, range='1Y')
            if all_ohlcvm is None:  # 该id有问题，无法获取数据
                print(f'{coins} 数据更新失败，获取出来为None')
                continue
            if coins in all_existing_coins_name:  # 数据更新拼接
                next_date = date_shifter(history_ohlcvm_latest_date, step='days', how_many=1)
                if all_ohlcvm.index[0] > next_date:  # 数据有遗漏，无法拼接
                    print(f'读取的最新数据跟历史数据接不上，请检查数据源是否有问题')
                    continue
                if all_ohlcvm.index[0] == history_ohlcvm.index[0]:  # 新数据可以直接覆盖旧数据
                    pass
                else:
                    all_ohlcvm = pd.concat([history_ohlcvm, all_ohlcvm.loc[next_date:, ]])
                    all_ohlcvm = all_ohlcvm[history_ohlcvm.columns]
                    all_ohlcvm.sort_index(inplace=True)
                    all_ohlcvm.index.name = 'end_date'

            all_ohlcvm.to_excel(f'{file_name}')  # 将新数据data写入coins工作表中
            print(f'{coins} 数据更新成功')
            time.sleep(3)  # 避免读取太频繁


def split_a_list_into_batches_and_deal_each_with_a_process(the_list, **kwargs):
    """
    将一个 list 分批，采用多进程并行处理每一批，并把所有的运行结果放入一个 list 返回

    :raise ValueError: 如果 what_version 不是 v1 或 v2
    :param the_list: 数据
    :param kwargs:
        - how_many，int，分成的批数
        - info，dict，传给单体处理函数的关键字参数
        - ignore_err，bool，是否忽略错误（针对个体运行时的），默认：False
        - shuffle_the_list，bool，是否打乱数组中元素的顺序，默认：False
        - func_2_run，function，数组分批后，如何处理每一批的函数，默认：None
        - what_version，str，使用函数的第几个版本的功能，默认：v1
    :return: list，包含了每一个个体被处理后的结果
    """
    how_many = kwargs.pop('how_many')
    info = kwargs.pop('info')
    ignore_err = kwargs.pop('ignore_err', False)
    shuffle_the_list = kwargs.pop('shuffle_the_list', False)
    func_2_run = kwargs.pop('func_2_run', None)
    what_version = kwargs.pop('what_version', 'v1')

    # 2019年12月22日14:22:21，如果这个数组为空的话怎么办？为了保持一致性
    # 因为这个函数的结果一般是用来进行循环的，所以返回一个空数组看起来是比较合理的
    if not the_list:
        print(f'传入的数组，False：{the_list}，为了保持结果的一致性，返回[]')
        return []

    if shuffle_the_list:
        random.shuffle(the_list)  # 在原list上进行操作，修改了原对象
    batches = split_a_list_into_batches(the_list, how_many)

    if func_2_run is None:
        if what_version == 'v1':
            pass
            # func_2_run = _deal_a_batch_with_a_loop
        elif what_version == 'v2':
            func_2_run = _deal_a_batch_with_a_loop_v2
        else:
            raise ValueError('what_version 不是 v1 或 v2')

    batch_res = batch_process(
        batches, func=func_2_run, info=info,
        ignore_err=ignore_err, multiple_threads=False, pool_size=how_many
    )

    res = []
    for a_piece in batch_res:
        res.extend(a_piece)
    return res



start_date = '2013-01-01'
initialize = False
gap_days = 2
end_date = get_yesterday(marker='with_n_dash')  # 计算截止昨日收盘
DATA_DIR = f'E:\\指标数据'
file_path = os.path.join(DATA_DIR, r'all_history_ohlcvm_coinmarketcap\all_history_ohlcvm_coinmarketcap')
all_filename_path_dict = get_all_filename_path(file_path)
all_existing_coins_name = [i.replace('.xlsx', '') for i in all_filename_path_dict.keys()]  # 之前已经保存过的一些数据

# 获取市场上所有币的信息
file_path_1 = os.path.join(DATA_DIR, 'all_history_ohlcvm_coinmarketcap')
file_name_1 = os.path.join(file_path_1, f'all_coin_list_coinmarketcap')
all_coin_list_df = pd.read_excel(f'{file_name_1}.xlsx', index_col='slug')
all_coin_list_df = all_coin_list_df[all_coin_list_df['status'] != 'untracked']  # 剔除尚未开始追踪的，因为获取不到数据
print(f'目前总共有{all_coin_list_df.shape[0]}个币的数据')

# 获取每个币的信息
id_series = all_coin_list_df['id']  # 所有币的id
id_series = pd.DataFrame(id_series)
id_series.reset_index(inplace=True)
id_series_list_for_test = []
split_num = 0
while split_num < len(id_series.index):
    if split_num + 10 > len(id_series.index):
        end_num = len(id_series.index)
    else:
        end_num = split_num + 10
    split_id_series = id_series.iloc[split_num:end_num, :].copy()
    split_id_series.set_index('slug', inplace=True)
    split_id_series = split_id_series['id']
    id_series_list_for_test.append(split_id_series)
    split_num = split_num + 10

# 批处理框架
info = {
    'func': fetch_one_coin_data,
    'file_path': file_path,
    'end_date': end_date,
    'initialize': initialize,
    'gap_days': gap_days,
    'all_existing_coins_name': all_existing_coins_name,
    'all_coin_list_df': all_coin_list_df,
    'all_filename_path_dict': all_filename_path_dict,
}

split_a_list_into_batches_and_deal_each_with_a_process(
    the_list=id_series_list_for_test, info=info, ignore_err=True, how_many=15, what_version='v2')  # 进程数不能太多，不然可能被封